import tkinter as tk
from tkinter import ttk, filedialog, messagebox, simpledialog
import cv2
import os
import shutil
from PIL import Image, ImageTk
from pathlib import Path
import json

class YOLOv5LabelingGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("YOLOv5 标注工具")
        self.root.geometry("1200x800")
        
        # 数据存储
        self.labels = {}
        self.current_image = None
        self.current_image_path = ""
        self.image_folder_path = ""
        self.image_save_path = ""
        self.label_save_path = ""
        self.current_boxes = []
        self.temp_point = None
        self.current_label_id = None
        self.image_files = []
        self.current_image_index = 0
        self.canvas_image = None
        self.scale_factor = 1.0
        
        # 颜色定义
        self.label_colors = [
            "#00FF00", "#FF0000", "#0000FF", "#FFFF00", "#FF00FF", 
            "#00FFFF", "#800080", "#FFA500", "#0080FF", "#80FF00"
        ]
        self.color_names = ["绿色", "红色", "蓝色", "黄色", "洋红", "青色", "紫色", "橙色", "浅蓝", "浅绿"]
        
        self.setup_ui()
        
    def setup_ui(self):
        """设置UI界面"""
        # 主框架
        main_frame = ttk.Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 左侧控制面板
        control_frame = ttk.Frame(main_frame, width=300)
        control_frame.pack(side=tk.LEFT, fill=tk.Y, padx=(0, 10))
        control_frame.pack_propagate(False)
        
        # 右侧图片显示区域
        image_frame = ttk.Frame(main_frame)
        image_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)
        
        # 设置控制面板
        self.setup_control_panel(control_frame)
        
        # 设置图片显示区域
        self.setup_image_panel(image_frame)
        
    def setup_control_panel(self, parent):
        """设置控制面板"""
        # 标题
        title_label = ttk.Label(parent, text="YOLOv5 标注工具", font=("Arial", 16, "bold"))
        title_label.pack(pady=(0, 20))
        
        # 路径设置区域
        path_frame = ttk.LabelFrame(parent, text="路径设置", padding=10)
        path_frame.pack(fill=tk.X, pady=(0, 10))
        
        ttk.Button(path_frame, text="选择图片文件夹", command=self.select_image_folder).pack(fill=tk.X, pady=2)
        ttk.Button(path_frame, text="选择图片保存路径", command=self.select_image_save_path).pack(fill=tk.X, pady=2)
        ttk.Button(path_frame, text="选择标签保存路径", command=self.select_label_save_path).pack(fill=tk.X, pady=2)
        
        # 标签设置区域
        label_frame = ttk.LabelFrame(parent, text="标签设置", padding=10)
        label_frame.pack(fill=tk.X, pady=(0, 10))
        
        ttk.Button(label_frame, text="添加标签", command=self.add_label).pack(fill=tk.X, pady=2)
        ttk.Button(label_frame, text="删除标签", command=self.delete_label).pack(fill=tk.X, pady=2)
        
        # 标签列表
        self.label_listbox = tk.Listbox(label_frame, height=6)
        self.label_listbox.pack(fill=tk.X, pady=5)
        self.label_listbox.bind('<<ListboxSelect>>', self.on_label_select)
        
        # 图片导航区域
        nav_frame = ttk.LabelFrame(parent, text="图片导航", padding=10)
        nav_frame.pack(fill=tk.X, pady=(0, 10))
        
        # 起始图片设置
        start_frame = ttk.Frame(nav_frame)
        start_frame.pack(fill=tk.X, pady=2)
        ttk.Label(start_frame, text="起始图片:").pack(side=tk.LEFT)
        self.start_entry = ttk.Entry(start_frame, width=10)
        self.start_entry.pack(side=tk.RIGHT)
        
        ttk.Button(nav_frame, text="开始标注", command=self.start_labeling).pack(fill=tk.X, pady=5)
        
        # 导航按钮
        nav_btn_frame = ttk.Frame(nav_frame)
        nav_btn_frame.pack(fill=tk.X, pady=2)
        ttk.Button(nav_btn_frame, text="上一张", command=self.prev_image).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 2))
        ttk.Button(nav_btn_frame, text="下一张", command=self.next_image).pack(side=tk.RIGHT, fill=tk.X, expand=True, padx=(2, 0))
        
        # 操作区域
        action_frame = ttk.LabelFrame(parent, text="操作", padding=10)
        action_frame.pack(fill=tk.X, pady=(0, 10))
        
        ttk.Button(action_frame, text="保存标注", command=self.save_current).pack(fill=tk.X, pady=2)
        ttk.Button(action_frame, text="清除标注", command=self.clear_annotations).pack(fill=tk.X, pady=2)
        ttk.Button(action_frame, text="保存并下一张", command=self.save_and_next).pack(fill=tk.X, pady=2)
        
        # 信息显示区域
        info_frame = ttk.LabelFrame(parent, text="信息", padding=10)
        info_frame.pack(fill=tk.BOTH, expand=True)
        
        self.info_text = tk.Text(info_frame, height=8, wrap=tk.WORD)
        scrollbar = ttk.Scrollbar(info_frame, orient=tk.VERTICAL, command=self.info_text.yview)
        self.info_text.configure(yscrollcommand=scrollbar.set)
        self.info_text.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        
    def setup_image_panel(self, parent):
        """设置图片显示面板"""
        # 图片信息标签
        self.image_info_label = ttk.Label(parent, text="请选择图片文件夹开始标注", font=("Arial", 12))
        self.image_info_label.pack(pady=(0, 10))
        
        # 图片画布
        canvas_frame = ttk.Frame(parent)
        canvas_frame.pack(fill=tk.BOTH, expand=True)
        
        self.canvas = tk.Canvas(canvas_frame, bg="white", cursor="crosshair")
        h_scrollbar = ttk.Scrollbar(canvas_frame, orient=tk.HORIZONTAL, command=self.canvas.xview)
        v_scrollbar = ttk.Scrollbar(canvas_frame, orient=tk.VERTICAL, command=self.canvas.yview)
        self.canvas.configure(xscrollcommand=h_scrollbar.set, yscrollcommand=v_scrollbar.set)
        
        self.canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        v_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        h_scrollbar.pack(side=tk.BOTTOM, fill=tk.X)
        
        # 绑定鼠标事件
        self.canvas.bind("<Button-1>", self.on_canvas_click)
        self.canvas.bind("<Motion>", self.on_canvas_motion)
        
    def log_info(self, message):
        """记录信息到文本框"""
        self.info_text.insert(tk.END, f"{message}\n")
        self.info_text.see(tk.END)
        self.root.update()
        
    def select_image_folder(self):
        """选择图片文件夹"""
        folder = filedialog.askdirectory(title="选择图片文件夹")
        if folder:
            self.image_folder_path = folder
            self.load_image_files()
            self.log_info(f"已选择图片文件夹: {folder}")
            
    def select_image_save_path(self):
        """选择图片保存路径"""
        folder = filedialog.askdirectory(title="选择图片保存路径")
        if folder:
            self.image_save_path = folder
            os.makedirs(folder, exist_ok=True)
            self.log_info(f"已选择图片保存路径: {folder}")
            
    def select_label_save_path(self):
        """选择标签保存路径"""
        folder = filedialog.askdirectory(title="选择标签保存路径")
        if folder:
            self.label_save_path = folder
            os.makedirs(folder, exist_ok=True)
            self.log_info(f"已选择标签保存路径: {folder}")
            
    def load_image_files(self):
        """加载图片文件列表"""
        if not self.image_folder_path:
            return
            
        valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
        image_files = []
        
        for file in os.listdir(self.image_folder_path):
            file_ext = os.path.splitext(file)[1].lower()
            if file_ext in valid_extensions:
                file_name = os.path.splitext(file)[0]
                try:
                    int(file_name)
                    image_files.append(file)
                except ValueError:
                    continue
        
        if image_files:
            # 按数字排序
            image_files.sort(key=lambda x: int(os.path.splitext(x)[0]))
            self.image_files = image_files
            self.log_info(f"找到 {len(image_files)} 个数字命名的图片文件")
            self.start_entry.delete(0, tk.END)
            self.start_entry.insert(0, "1")
        else:
            messagebox.showwarning("警告", "未找到有效的数字命名图片文件")
            
    def add_label(self):
        """添加标签"""
        dialog = tk.Toplevel(self.root)
        dialog.title("添加标签")
        dialog.geometry("300x150")
        dialog.transient(self.root)
        dialog.grab_set()
        
        ttk.Label(dialog, text="标签ID:").pack(pady=5)
        id_entry = ttk.Entry(dialog)
        id_entry.pack(pady=5)
        
        ttk.Label(dialog, text="标签名称:").pack(pady=5)
        name_entry = ttk.Entry(dialog)
        name_entry.pack(pady=5)
        
        def confirm():
            try:
                label_id = int(id_entry.get())
                label_name = name_entry.get().strip()
                if label_name:
                    self.labels[label_id] = label_name
                    self.update_label_listbox()
                    color_name = self.color_names[label_id % len(self.color_names)]
                    self.log_info(f"已添加标签: {label_id} -> {label_name} ({color_name})")
                    dialog.destroy()
                else:
                    messagebox.showwarning("警告", "请输入标签名称")
            except ValueError:
                messagebox.showwarning("警告", "标签ID必须是数字")
        
        ttk.Button(dialog, text="确定", command=confirm).pack(pady=10)
        
    def delete_label(self):
        """删除标签"""
        selection = self.label_listbox.curselection()
        if selection:
            index = selection[0]
            label_items = list(self.labels.items())
            if index < len(label_items):
                label_id, label_name = label_items[index]
                del self.labels[label_id]
                self.update_label_listbox()
                self.log_info(f"已删除标签: {label_id} -> {label_name}")
        else:
            messagebox.showwarning("警告", "请选择要删除的标签")
            
    def update_label_listbox(self):
        """更新标签列表框"""
        self.label_listbox.delete(0, tk.END)
        for label_id, label_name in sorted(self.labels.items()):
            color_name = self.color_names[label_id % len(self.color_names)]
            self.label_listbox.insert(tk.END, f"{label_id}: {label_name} ({color_name})")
            
    def on_label_select(self, event):
        """标签选择事件"""
        selection = self.label_listbox.curselection()
        if selection:
            index = selection[0]
            label_items = list(self.labels.items())
            if index < len(label_items):
                self.current_label_id = label_items[index][0]
                label_name = label_items[index][1]
                color_name = self.color_names[self.current_label_id % len(self.color_names)]
                self.log_info(f"已选择标签: {self.current_label_id} ({label_name}, {color_name})")
                
    def start_labeling(self):
        """开始标注"""
        if not self.image_files:
            messagebox.showwarning("警告", "请先选择图片文件夹")
            return
        if not self.labels:
            messagebox.showwarning("警告", "请先添加标签")
            return
        if not self.image_save_path or not self.label_save_path:
            messagebox.showwarning("警告", "请设置图片和标签保存路径")
            return
            
        try:
            start_num = int(self.start_entry.get())
            if 1 <= start_num <= len(self.image_files):
                self.current_image_index = start_num - 1
                self.load_current_image()
                self.log_info(f"开始标注，从第{start_num}张图片开始")
            else:
                messagebox.showwarning("警告", f"起始图片编号应在1-{len(self.image_files)}之间")
        except ValueError:
            messagebox.showwarning("警告", "请输入有效的起始图片编号")
            
    def load_current_image(self):
        """加载当前图片"""
        if self.current_image_index >= len(self.image_files):
            messagebox.showinfo("完成", "所有图片标注完成！")
            return False
            
        current_file = self.image_files[self.current_image_index]
        image_path = os.path.join(self.image_folder_path, current_file)
        
        # 使用OpenCV读取图片
        cv_image = cv2.imread(image_path)
        if cv_image is not None:
            # 转换为RGB格式
            cv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
            self.current_image = cv_image
            self.current_image_path = image_path
            
            # 转换为PIL图片并显示
            pil_image = Image.fromarray(cv_image)
            
            # 计算缩放比例以适应画布
            canvas_width = self.canvas.winfo_width()
            canvas_height = self.canvas.winfo_height()
            if canvas_width > 1 and canvas_height > 1:
                img_width, img_height = pil_image.size
                scale_x = canvas_width / img_width
                scale_y = canvas_height / img_height
                self.scale_factor = min(scale_x, scale_y, 1.0)  # 不放大，只缩小
                
                new_width = int(img_width * self.scale_factor)
                new_height = int(img_height * self.scale_factor)
                pil_image = pil_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
            
            self.canvas_image = ImageTk.PhotoImage(pil_image)
            self.canvas.delete("all")
            self.canvas.create_image(0, 0, anchor=tk.NW, image=self.canvas_image)
            self.canvas.configure(scrollregion=self.canvas.bbox("all"))
            
            # 重置标注
            self.current_boxes = []
            self.temp_point = None
            
            # 更新信息
            self.image_info_label.config(text=f"图片: {current_file} (第{self.current_image_index + 1}/{len(self.image_files)}张)")
            self.log_info(f"已加载: {current_file}")
            
            # 重绘标注框
            self.draw_annotations()
            return True
        else:
            messagebox.showerror("错误", f"无法加载图片: {current_file}")
            return False

    def on_canvas_click(self, event):
        """画布点击事件"""
        if self.current_label_id is None:
            messagebox.showwarning("警告", "请先选择标签")
            return
        if self.current_image is None:
            return

        # 获取画布坐标
        canvas_x = self.canvas.canvasx(event.x)
        canvas_y = self.canvas.canvasy(event.y)

        # 转换为原图坐标
        orig_x = int(canvas_x / self.scale_factor)
        orig_y = int(canvas_y / self.scale_factor)

        if self.temp_point is None:
            self.temp_point = (orig_x, orig_y)
            self.log_info(f"左上角点: ({orig_x}, {orig_y})")
            # 在画布上绘制临时点
            canvas_point_x = orig_x * self.scale_factor
            canvas_point_y = orig_y * self.scale_factor
            self.canvas.create_oval(canvas_point_x-5, canvas_point_y-5, canvas_point_x+5, canvas_point_y+5,
                                  fill="red", outline="red", tags="temp_point")
        else:
            x1, y1 = self.temp_point
            left, top = min(x1, orig_x), min(y1, orig_y)
            right, bottom = max(x1, orig_x), max(y1, orig_y)

            # 检查矩形大小
            if abs(right - left) < 10 or abs(bottom - top) < 10:
                messagebox.showwarning("警告", "矩形太小，请重新选择")
                self.temp_point = None
                self.canvas.delete("temp_point")
                return

            # 转换为YOLO格式
            img_h, img_w = self.current_image.shape[:2]
            center_x = (left + right) / 2.0 / img_w
            center_y = (top + bottom) / 2.0 / img_h
            width = (right - left) / img_w
            height = (bottom - top) / img_h

            box_info = {
                'label_id': self.current_label_id,
                'center_x': center_x,
                'center_y': center_y,
                'width': width,
                'height': height,
                'pixel_coords': (left, top, right, bottom)
            }
            self.current_boxes.append(box_info)

            label_name = self.labels[self.current_label_id]
            color_name = self.color_names[self.current_label_id % len(self.color_names)]
            self.log_info(f"添加标注框: 标签{self.current_label_id}({label_name},{color_name}) 坐标({left},{top},{right},{bottom})")

            self.temp_point = None
            self.canvas.delete("temp_point")
            self.draw_annotations()

    def on_canvas_motion(self, event):
        """鼠标移动事件"""
        if self.temp_point is not None and self.current_image is not None:
            # 清除之前的预览线
            self.canvas.delete("preview_rect")

            # 获取当前鼠标位置
            canvas_x = self.canvas.canvasx(event.x)
            canvas_y = self.canvas.canvasy(event.y)

            # 获取起始点的画布坐标
            start_canvas_x = self.temp_point[0] * self.scale_factor
            start_canvas_y = self.temp_point[1] * self.scale_factor

            # 绘制预览矩形
            self.canvas.create_rectangle(start_canvas_x, start_canvas_y, canvas_x, canvas_y,
                                       outline="red", width=2, tags="preview_rect")

    def draw_annotations(self):
        """绘制标注框"""
        # 清除之前的标注
        self.canvas.delete("annotation")

        for i, box in enumerate(self.current_boxes):
            left, top, right, bottom = box['pixel_coords']
            label_id = box['label_id']
            label_name = self.labels[label_id]

            # 转换为画布坐标
            canvas_left = left * self.scale_factor
            canvas_top = top * self.scale_factor
            canvas_right = right * self.scale_factor
            canvas_bottom = bottom * self.scale_factor

            # 获取颜色
            color = self.label_colors[label_id % len(self.label_colors)]

            # 绘制矩形框
            self.canvas.create_rectangle(canvas_left, canvas_top, canvas_right, canvas_bottom,
                                       outline=color, width=3, tags="annotation")

            # 绘制标签文本背景
            label_text = f"{label_id}: {label_name}"
            text_id = self.canvas.create_text(canvas_left + 5, canvas_top - 10, text=label_text,
                                            anchor=tk.NW, fill=color, font=("Arial", 10, "bold"),
                                            tags="annotation")

            # 获取文本边界框并绘制背景
            bbox = self.canvas.bbox(text_id)
            if bbox:
                self.canvas.create_rectangle(bbox[0]-2, bbox[1]-2, bbox[2]+2, bbox[3]+2,
                                           fill="white", outline=color, tags="annotation")
                # 重新绘制文本在背景之上
                self.canvas.create_text(canvas_left + 5, canvas_top - 10, text=label_text,
                                      anchor=tk.NW, fill=color, font=("Arial", 10, "bold"),
                                      tags="annotation")

    def prev_image(self):
        """上一张图片"""
        if self.current_image_index > 0:
            self.current_image_index -= 1
            self.load_current_image()
        else:
            messagebox.showinfo("提示", "已经是第一张图片")

    def next_image(self):
        """下一张图片"""
        if self.current_image_index < len(self.image_files) - 1:
            self.current_image_index += 1
            self.load_current_image()
        else:
            messagebox.showinfo("提示", "已经是最后一张图片")

    def save_current(self):
        """保存当前标注"""
        if not self.current_image_path:
            return

        # 保存标签文件
        image_name = Path(self.current_image_path).stem
        label_file_path = os.path.join(self.label_save_path, f"{image_name}.txt")

        try:
            with open(label_file_path, 'w') as f:
                for box in self.current_boxes:
                    line = f"{box['label_id']} {box['center_x']:.6f} {box['center_y']:.6f} {box['width']:.6f} {box['height']:.6f}\n"
                    f.write(line)
            self.log_info(f"已保存标签文件: {image_name}.txt (包含{len(self.current_boxes)}个标注框)")
        except Exception as e:
            messagebox.showerror("错误", f"保存标签文件失败: {e}")
            return

        # 复制图片文件
        try:
            image_name_full = os.path.basename(self.current_image_path)
            dest_path = os.path.join(self.image_save_path, image_name_full)
            if self.current_image_path != dest_path:
                shutil.copy2(self.current_image_path, dest_path)
                self.log_info(f"已复制图片: {image_name_full}")
        except Exception as e:
            messagebox.showerror("错误", f"复制图片失败: {e}")

    def clear_annotations(self):
        """清除当前标注"""
        self.current_boxes = []
        self.temp_point = None
        self.canvas.delete("annotation")
        self.canvas.delete("temp_point")
        self.canvas.delete("preview_rect")
        self.log_info("已清除所有标注")

    def save_and_next(self):
        """保存并下一张"""
        self.save_current()
        self.next_image()

def main():
    root = tk.Tk()
    app = YOLOv5LabelingGUI(root)
    root.mainloop()

if __name__ == "__main__":
    main()
